# Lint as: python3
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Market data config description.

Market data configuration has a form

market_data_config = {
    "Currency":  {
        risk_free_curve : RateConfig(RateCurveOption.BOOTSTRAP,
                                     ["list of bootstrap instruments"]),
        rate_index : RateConfig(RateCurveOption.BOOTSTRAP,
                                ["list of bootstrap instruments"]),
        surface_id: "to be specified",
        fixings: "to be specified"},
    "Asset": "to be specified"}
"""

import enum
from typing import List, Union

import dataclasses

from tf_quant_finance.experimental.pricing_platform.framework.core import daycount_conventions
from tf_quant_finance.experimental.pricing_platform.framework.core import interpolation_method
from tf_quant_finance.experimental.pricing_platform.framework.core import types

_DayCountConventions = daycount_conventions.DayCountConventions
_InterpolationMethod = interpolation_method.InterpolationMethod

# Can't infer autogenerated type of DayCountConventions so extracting
# explicitly
_DayCountConventionsProtoType = types.DayCountConventionsProtoType


class RateBootstrapOption(enum.Enum):
  FROM_CSV = "FROM_CSV"
  FROM_JSON = "FROM_JSON"
  FROM_BQ = "FROM_BQ"  # From BigQuery
  BOOTSTRAP = "BOOTSTRAP"


class RateVolBootstrapOption(enum.Enum):
  FROM_CSV = "FROM_CSV"
  FROM_JSON = "FROM_JSON"
  FROM_BQ = "FROM_BQ"  # From BigQuery
  BOOTSTRAP = "BOOTSTRAP"


@dataclasses.dataclass
class RateConfig:
  file_id: Union[str, None] = None
  rate_bootstrap_option: Union[RateBootstrapOption, None] = None
  bootstrap_instrumetns: Union[List[str], None] = None
  interpolation_method: _InterpolationMethod = _InterpolationMethod.CUBIC
  interpolate_rates: bool = True
  daycount_convention: _DayCountConventionsProtoType = _DayCountConventions.ACTUAL_365


__all__ = ["RateBootstrapOption",
           "RateVolBootstrapOption",
           "RateConfig"]
